import math
from typing import Callable
import pysam
import pandas as pd
import numpy as np

from multiprocessing import Pool


def calculate_coverage(bam: pysam.AlignmentFile, chrom: str, start: int, end: int) -> dict:
    count_coverage = np.array(bam.count_coverage(chrom, start, end)).sum(axis=0)
    min_depth = count_coverage.min()
    max_depth = count_coverage.max()
    avg_depth = count_coverage.mean()
    cov_1x = len(count_coverage[count_coverage > 0])/(end-start)
    cov_4x = len(count_coverage[count_coverage >= 4])/(end-start)
    cov_10x = len(count_coverage[count_coverage >= 10])/(end-start)
    cov_20x = len(count_coverage[count_coverage >= 20])/(end-start)
    cov_50x = len(count_coverage[count_coverage >= 50])/(end-start)
    cov_100x = len(count_coverage[count_coverage >= 100])/(end-start)
    cov_200x = len(count_coverage[count_coverage >= 200])/(end-start)
    return {
        'Chromosome': chrom,
        'Start': start,
        'End': end,
        'Min Depth': min_depth,
        'Max Depth': max_depth,
        'Avg Depth': avg_depth,
        'Coverage(>=1X)': cov_1x,
        'Coverage(>=4X)': cov_4x,
        'Coverage(>=10X)': cov_10x,
        'Coverage(>=20X)': cov_20x,
        'Coverage(>=50X)': cov_50x,
        'Coverage(>=100X)': cov_100x,
        'Coverage(>=200X)': cov_200x,
    }


def calculate_transcript_coverage(bam_file: str, gpes: pd.DataFrame) -> list[dict]:
    bam = pysam.AlignmentFile(bam_file, "r")
    if not bam.check_index():
        raise Exception("ERROR: bam index not found")
    rows = list()
    for gpe in gpes.iloc:
        starts = [int(p) for p in gpe.ExonStarts.rstrip(',').split(',')]
        ends = [int(p) for p in gpe.ExonEnds.rstrip(',').split(',')]
        for i in range(int(gpe.ExonCount)):
            exon = f'Exon{(i+1) if gpe.Strand == "+" else (gpe.ExonCount-i)}'
            row = {'Gene': gpe.Gene, 'Transcript': gpe.Transcript, 'Exon': exon}
            row.update(calculate_coverage(bam, str(gpe.Chrom), starts[i], ends[i]))
            rows.append(row)
    return rows


def calculate_bed_coverage(bam_file: str, beds: pd.DataFrame) -> list[dict]:
    bam = pysam.AlignmentFile(bam_file, "r")
    if not bam.check_index():
        raise Exception("ERROR: bam index not found")
    rows = list()
    for bed in beds.iloc:
        row = {"Name": bed.Name}
        row.update(calculate_coverage(bam, str(bed.Chrom), int(bed.End), int(bed.End)))
        rows.append(row)
    return rows


def multiple_calculate_coverages(bam_file: str, dataframe: pd.DataFrame, thread: int, cal_func: Callable[[str, pd.DataFrame], list]) -> pd.DataFrame:
    dataframe = dataframe.head(10)
    bin_num = 100
    bin_size = math.ceil(dataframe.shape[0]/bin_num)
    pool = Pool(thread)
    jobs = [pool.apply_async(cal_func, args=(bam_file, dataframe.loc[i*bin_size:(i+1)*bin_size-1])) for i in range(bin_num)]
    pool.close()
    pool.join()
    results = list()
    for job in jobs:
        results.extend(job.get())
    return pd.DataFrame(results)


def calculate_transcript_coverages(bam_file: str, gpes: pd.DataFrame, thread: int) -> pd.DataFrame:
    return multiple_calculate_coverages(bam_file, gpes, thread, calculate_transcript_coverage)


def calculate_bed_coverages(bam_file: str, beds: pd.DataFrame, thread: int) -> pd.DataFrame:
    return multiple_calculate_coverages(bam_file, beds, thread, calculate_bed_coverage)
